# coding=utf-8

import hiai
from hiai.nn_tensor_lib import DataType
from utils import jpegHandler
import os
import numpy as np
import time
import cv2 as cv

class_names = ('_bg_',
    'aeroplane', 'bicycle', 'bird', 'boat',
    'bottle', 'bus', 'car', 'cat', 'chair',
    'cow', 'diningtable', 'dog', 'horse',
    'motorbike', 'person', 'pottedplant',
    'sheep', 'sofa', 'train', 'tvmonitor')
num_class = len(class_names)
index_to_class = dict(zip(range(num_class), class_names))

omFileName = './models/res101_improve.om'
srcFileDir = './JPEGImages/'
dstFileDir = './Result/rfcn_VOC/'
kCategoryIndex = 2

def py_cpu_nms(dets, thresh):
    """Pure Python NMS baseline."""
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1]

    return keep

def CreateGraph(model, modelInWidth, modelInHeight, dvppInWidth, dvppInHeight):
    myGraph = hiai.hiai._global_default_graph_stack.get_default_graph()
    if myGraph is None:
        print 'get defaule graph failed'
        return None
    print 'dvppwidth %d, dvppheight %d' % (dvppInWidth, dvppInHeight)

    cropConfig = hiai.CropConfig(0, 0, dvppInWidth, dvppInHeight)
    print 'cropConfig ', cropConfig
    resizeConfig = hiai.ResizeConfig(modelInWidth, modelInHeight)
    print 'resizeConfig ', resizeConfig

    nntensorList = hiai.NNTensorList()
    print 'nntensorList', nntensorList

    resultCrop = hiai.crop(nntensorList, cropConfig)
    print 'resultCrop', resultCrop

    resultResize = hiai.resize(resultCrop, resizeConfig)
    print 'resultResize', resultResize

    resultInference = hiai.inference(resultResize, model, None)
    print 'resultInference', resultInference

    if (hiai.HiaiPythonStatust.HIAI_PYTHON_OK == myGraph.create_graph()):
        print 'create graph ok !!!!'
        return myGraph
    else:
        print 'create graph failed, please check Davinc log.'
        return None


def CreateGraphWithoutDVPP(model):
    print model
    myGraph = hiai.hiai._global_default_graph_stack.get_default_graph()
    print myGraph
    if myGraph is None:
        print 'get defaule graph failed'
        return None

    nntensorList = hiai.NNTensorList()
    print nntensorList

    resultInference = hiai.inference(nntensorList, model, None)
    print nntensorList
    print hiai.HiaiPythonStatust.HIAI_PYTHON_OK
    # print myGraph.create_graph()

    if (hiai.HiaiPythonStatust.HIAI_PYTHON_OK == myGraph.create_graph()):
        print 'create graph ok !!!!'
        return myGraph
    else:
        print 'create graph failed, please check Davinc log.'
        return None


def GraphInference(graphHandle, inputTensorList):
    if not isinstance(graphHandle, hiai.Graph):
        print "graphHandle is not Graph object"
        return None

    resultList = graphHandle.proc(inputTensorList)
    return resultList


def PostProcess(resultList, dstFilePath, fileName):
    if resultList is not None:
        tensor_num = np.reshape(resultList[0], (304,21))
        tensor_bbox = np.reshape(resultList[1], (304,8))
        tensor_rois = np.reshape(resultList[2], (304,13))
        tensor_score = np.argsort(tensor_num,axis =1)[:,20]
        img = cv.imread(fileName)
        img_rows, img_cols, img_channel = img.shape
        CONF_THRESH = 0.8
        NMS_THRESH = 0.3
        bboxs = np.zeros((304,4))
        bboxs[:, 0] = img_cols*tensor_rois[:, 0] + tensor_rois[:, 9]
        bboxs[:, 1] = img_rows*tensor_rois[:, 1] + tensor_rois[:, 10]
        bboxs[:, 2] = img_cols*tensor_rois[:, 2] + tensor_rois[:, 11]
        bboxs[:, 3] = img_rows*tensor_rois[:, 3] + tensor_rois[:, 12]
        for j in range(1,len(class_names)):
            scores = tensor_num[:,j]
            boxes = np.hstack((bboxs, scores[:, np.newaxis].astype(np.float32)))
            keep_index = py_cpu_nms(boxes, NMS_THRESH)
            boxes_nms = boxes[keep_index,:]
            keep_final = np.where(boxes_nms[:, -1] >= CONF_THRESH)[0]
            if len(keep_final)==0:
                continue
            for i in keep_final:
                lt_x = int(boxes_nms[i][0])
                lt_y = int(boxes_nms[i][1])
                # rb_x = int(boxes_nms[i][0]+boxes_nms[i][2])
                # rb_y = int(boxes_nms[i][1]+boxes_nms[i][3])
                rb_x = int(boxes_nms[i][2])
                rb_y = int(boxes_nms[i][3])
                score = boxes_nms[i][4]
                cv.rectangle(img, (lt_x, lt_y ), (rb_x, rb_y), (0, 0, 255), 4)
                cv.putText(img, "class:%s,score:%.3f" % (index_to_class[j], score), (lt_x, lt_y),
                           cv.FONT_HERSHEY_COMPLEX, 0.5, (0, 255, 255), 1)

        _, filename = os.path.split(fileName)
        cv.imwrite("%s%s" % (dstFilePath, filename), img)
        return None
    else:
        print 'graph inference failed '
        return None



def main():
    # capture_flag = 1  # use camera or not
    inferenceModel = hiai.AIModelDescription('faster-rcnn', omFileName)
    print omFileName
    print inferenceModel
    # we will resize the jpeg to 896*608 to meet faster-rcnn requirement via opencv,
    # so DVPP resizing is not needed
    myGraph = CreateGraphWithoutDVPP(inferenceModel)
    if myGraph is None:
        print "CreateGraph failed"
        return None

    # in this sample demo, the faster-rcnn  model requires 896*608 images
    dvppInWidth = 256
    dvppInHeight = 224

    start = time.time()
    index = 0

    pathDir = sorted(os.listdir(srcFileDir))
    for allDir in pathDir:
        child = os.path.join('%s%s' % (srcFileDir, allDir))
        if (not jpegHandler.is_img(child)):
            print '[info] file : ' + child + ' is not image !'
            continue

        # read the jpeg file and resize it to required w&h, than change it to YUV format.
        input_image = jpegHandler.jpeg2yuv(child, dvppInWidth, dvppInHeight)

        inputImageTensor = hiai.NNTensor(input_image, dvppInWidth, dvppInHeight, 3, 'testImage', DataType.UINT8_T,
                                         dvppInWidth * dvppInHeight * 3 / 2)
        imageinfo = np.array([dvppInWidth, dvppInHeight, 3]).astype(np.float32)
        imageinfo = np.reshape(imageinfo, (1, 3))
        infoTensor = hiai.NNTensor(imageinfo, 1, 3, 1, 'testinfo', DataType.FLOAT32_T, imageinfo.size)

        datalist = [inputImageTensor, infoTensor]
        nntensorList = hiai.NNTensorList(datalist)
        if not nntensorList:
            print "nntensorList is null"
            break
        resultList = GraphInference(myGraph, nntensorList)

        if resultList is None:
            print "graph inference failed"
            continue
        print resultList[1].shape
        PostProcess(resultList, dstFileDir, child)

    end = time.time()
    print 'cost time ' + str((end - start) * 1000) + 'ms'

    hiai.hiai._global_default_graph_stack.get_default_graph().destroy()

    print '-------------------end'


if __name__ == "__main__":
    main()
